package com.campus.counseling.service.impl;

import com.campus.counseling.config.LSTMConfig;
import com.campus.counseling.entity.TrainingData;
import com.campus.counseling.model.mapper.ChatAnalysisMapper;
import com.campus.counseling.model.mapper.ChatMessageMapper;
import com.campus.counseling.service.TrainingDataService;
import com.campus.counseling.service.Word2VecService;
import com.campus.counseling.service.DL4JModelService;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.nd4j.linalg.api.ndarray.INDArray;
import org.nd4j.linalg.dataset.DataSet;
import org.nd4j.linalg.factory.Nd4j;
import org.nd4j.linalg.indexing.INDArrayIndex;
import org.springframework.stereotype.Service;
import java.util.List;
import java.util.stream.Collectors;
import org.nd4j.linalg.indexing.NDArrayIndex;
import java.util.Arrays;

@Slf4j
@Service
@RequiredArgsConstructor
public class TrainingDataServiceImpl implements TrainingDataService {

    // 词向量服务
    private final Word2VecService word2VecService;
    // LSTM配置
    private final LSTMConfig lstmConfig;
    private final DL4JModelService dl4jModelService;

    @Override
    public DataSet prepareTrainingData(List<TrainingData> trainingData) {
        try {
            log.info("准备训练数据，数据量: {}", trainingData.size());
            
            // 1. 提取文本和情感分数
            List<String> texts = trainingData.stream()
                .map(data -> data.getQuestion() + " " + data.getDescription())
                .collect(Collectors.toList());
                
            List<Double> scores = trainingData.stream()
                .map(TrainingData::getEmotionScore)
                .collect(Collectors.toList());
            
            // 2. 提取特征
            INDArray features = dl4jModelService.extractFeatures(texts);
            
            // 3. 创建标签矩阵 [batchSize, timeSteps, 1]
            INDArray labels = Nd4j.create(scores.size(), 1, 1);  // 修改为3维
            for (int i = 0; i < scores.size(); i++) {
                labels.putScalar(new int[]{i, 0, 0}, scores.get(i));  // 使用3维索引
            }
            
            // 4. 记录维度信息
            log.info("特征维度: {}, 标签维度: {}", 
                Arrays.toString(features.shape()),
                Arrays.toString(labels.shape()));
            
            return new DataSet(features, labels);
            
        } catch (Exception e) {
            log.error("准备训练数据失败: ", e);
            throw new RuntimeException("准备训练数据失败: " + e.getMessage());
        }
    }

    private double[] extractEmotionFeatures(String text) {
        // 简单的情感词统计
        double[] features = new double[5];  // 固定5个情感特征维度
        
        // 积极情感词
        String[] positiveWords = {"开心", "希望", "充实", "愉快", "好"};
        // 消极情感词
        String[] negativeWords = {"难过", "痛苦", "绝望", "失败", "死"};
        // 程度词
        String[] degreeWords = {"很", "非常", "特别"};
        // 否定词
        String[] negationWords = {"不", "没", "别"};
        
        // 统计词频
        for (String word : positiveWords) {
            if (text.contains(word)) {
                features[0] += 1.0;
            }
        }
        for (String word : negativeWords) {
            if (text.contains(word)) {
                features[1] += 1.0;
            }
        }
        for (String word : degreeWords) {
            if (text.contains(word)) {
                features[2] += 1.0;
            }
        }
        for (String word : negationWords) {
            if (text.contains(word)) {
                features[3] += 1.0;
            }
        }
        
        // 计算文本长度特征
        features[4] = text.length() / 100.0;  // 归一化文本长度
        
        return features;
    }

} 